{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "659a189e-33c9-426f-b318-7cb2f433ae4a",
   "metadata": {
    "id": "659a189e-33c9-426f-b318-7cb2f433ae4a"
   },
   "source": [
    "<center>\n",
    "<h4>CDS 110, Lecture 8b</h4>\n",
    "<font color=blue><h1>Full Controller Stack for a Planar Vertical Take-Off and Landing (PVTOL) System</h1></font>\n",
    "<h3>Richard M. Murray, Winter 2024</h4>\n",
    "</center>\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/drive/1XulsQqbthMkr3g58OTctIYKYpqirOgns)\n",
    "\n",
    "The purpose of this lecture is to introduce tools that can be used for frequency domain modeling and analysis of linear systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1be7545a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "from math import sin, cos, pi\n",
    "from scipy.optimize import NonlinearConstraint\n",
    "import time\n",
    "try:\n",
    "  import control as ct\n",
    "  print(\"python-control\", ct.__version__)\n",
    "except ImportError:\n",
    "  !pip install control\n",
    "  import control as ct\n",
    "import control.optimal as opt\n",
    "import control.flatsys as fs\n",
    "\n",
    "# Use control parameters for plotting\n",
    "plt.rcParams.update(ct.rcParams)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5a1858a",
   "metadata": {
    "id": "c5a1858a"
   },
   "source": [
    "## System definition\n",
    "\n",
    "Consider the PVTOL system `pvtol_noisy`, defined in `pvtol.py`:\n",
    "\n",
    "$$\n",
    "  \\begin{aligned}\n",
    "    m \\ddot x &= F_1 \\cos\\theta - F_2 \\sin\\theta - c \\dot x + D_x, \\\\\n",
    "    m \\ddot y &= F_1 \\sin\\theta + F_2 \\cos\\theta - c \\dot y - m g + D_y, \\\\\n",
    "    J \\ddot \\theta &= r F_1,\n",
    "  \\end{aligned} \\qquad\n",
    "  \\vec Y =\n",
    "    \\begin{bmatrix} x \\\\ y \\\\ \\theta \\end{bmatrix} +\n",
    "    \\begin{bmatrix} N_x \\\\ N_y \\\\ N_z \\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "Assume that the input disturbances are modeled by independent, first\n",
    "order Markov (Ornstein-Uhlenbeck) processes with\n",
    "$Q_D = \\text{diag}(0.01, 0.01)$ and $\\omega_0 = 1$ and that the noise\n",
    "is modeled as white noise with covariance matrix\n",
    "\n",
    "$$\n",
    "  Q_N = \\begin{bmatrix}\n",
    "    2 \\times 10^{-4} & 0 & 1 \\times 10^{-5} \\\\\n",
    "    0 & 2 \\times 10^{-4} & 1 \\times 10^{-5} \\\\\n",
    "    1 \\times 10^{-5} & 1 \\times 10^{-5} & 1 \\times 10^{-4}\n",
    "  \\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "We will design a controller consisting of a trajectory generation module, a\n",
    "gain-scheduled, trajectory tracking module, and a state estimation\n",
    "module the moves the system from the origin to the equilibrum point\n",
    "point $x_\\text{f}$, $y_\\text{f}$ = 10, 0 while satisfying the\n",
    "constraint $0.5 \\sin(\\pi x / 10) - 0.1 \\leq y \\leq 1$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "D1aFeNuglL4a",
   "metadata": {
    "id": "D1aFeNuglL4a"
   },
   "source": [
    "We start by creating the PVTOL system without noise or disturbances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c32ec3f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# STANDARD PVTOL DYNAMICS\n",
    "def _pvtol_update(t, x, u, params):\n",
    "\n",
    "    # Get the parameter values\n",
    "    m = params.get('m', 4.)             # mass of aircraft\n",
    "    J = params.get('J', 0.0475)         # inertia around pitch axis\n",
    "    r = params.get('r', 0.25)           # distance to center of force\n",
    "    g = params.get('g', 9.8)            # gravitational constant\n",
    "    c = params.get('c', 0.05)           # damping factor (estimated)\n",
    "\n",
    "    # Get the inputs and states\n",
    "    x, y, theta, xdot, ydot, thetadot = x\n",
    "    F1, F2 = u\n",
    "\n",
    "    # Constrain the inputs\n",
    "    F2 = np.clip(F2, 0, 1.5 * m * g)\n",
    "    F1 = np.clip(F1, -0.1 * F2, 0.1 * F2)\n",
    "\n",
    "    # Dynamics\n",
    "    xddot = (F1 * cos(theta) - F2 * sin(theta) - c * xdot) / m\n",
    "    yddot = (F1 * sin(theta) + F2 * cos(theta) - m * g - c * ydot) / m\n",
    "    thddot = (r * F1) / J\n",
    "\n",
    "    return np.array([xdot, ydot, thetadot, xddot, yddot, thddot])\n",
    "\n",
    "# Define pvtol output function to only be x, y, and theta\n",
    "def _pvtol_output(t, x, u, params):\n",
    "    return x[0:3]\n",
    "\n",
    "# Create nonlinear input-output system of nominal pvtol system\n",
    "pvtol_nominal = ct.nlsys(\n",
    "    _pvtol_update, _pvtol_output, name=\"pvtol_nominal\",\n",
    "    states = [f'x{i}' for i in range(6)],\n",
    "    inputs = ['F1', 'F2'],\n",
    "    outputs = [f'x{i}' for i in range(3)]\n",
    ")\n",
    "\n",
    "print(pvtol_nominal)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "TTMQAAhFldW7",
   "metadata": {
    "id": "TTMQAAhFldW7"
   },
   "source": [
    "Next, we create a PVTOL system with noise and disturbances. This system will use the nominal PVTOL system and add disturbances as inputs to the state dynamics and noise to the system output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tqSvuzvOkps1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add wind and noise to system dynamics\n",
    "def _noisy_update(t, x, u, params):\n",
    "    # Get the inputs\n",
    "    F1, F2, Dx, Dy = u[:4]\n",
    "    if u.shape[0] > 4:\n",
    "        Nx, Ny, Nth = u[4:]\n",
    "    else:\n",
    "        Nx, Ny, Nth = 0, 0, 0\n",
    "\n",
    "    # Get the system response from the original dynamics\n",
    "    xdot, ydot, thetadot, xddot, yddot, thddot = \\\n",
    "        _pvtol_update(t, x, [F1, F2], params)\n",
    "\n",
    "    # Get the parameter values we need\n",
    "    m = params.get('m', 4.)             # mass of aircraft\n",
    "    J = params.get('J', 0.0475)         # inertia around pitch axis\n",
    "\n",
    "    # Now add the disturbances\n",
    "    xddot += Dx / m\n",
    "    yddot += Dy / m\n",
    "\n",
    "    return np.array([xdot, ydot, thetadot, xddot, yddot, thddot])\n",
    "\n",
    "# Define pvtol_noisy output function to only be x, y, and theta\n",
    "def _noisy_output(t, x, u, params):\n",
    "    F1, F2, Dx, Dy, Nx, Ny, Nth = u\n",
    "    return x[0:3] + np.array([Nx, Ny, Nth])\n",
    "\n",
    "# CREATE NONLINEAR INPUT-OUTPUT SYSTEM\n",
    "pvtol_noisy = ct.nlsys(\n",
    "    _noisy_update, _noisy_output, name=\"pvtol_noisy\",\n",
    "    states = [f'x{i}' for i in range(6)],\n",
    "    inputs = ['F1', 'F2'] + ['Dx', 'Dy'] + ['Nx', 'Ny', 'Nth'],\n",
    "    outputs = ['x', 'y', 'theta'],\n",
    "    params = {\n",
    "        'm': 4.,                # mass of aircraft\n",
    "        'J': 0.0475,            # inertia around pitch axis\n",
    "        'r': 0.25,              # distance to center of force\n",
    "        'g': 9.8,               # gravitational constant\n",
    "        'c': 0.05,              # damping factor (estimated)\n",
    "    }\n",
    ")\n",
    "\n",
    "print(pvtol_noisy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "057cba8f-79bd-4a45-a184-2424c569785d",
   "metadata": {
    "id": "057cba8f-79bd-4a45-a184-2424c569785d"
   },
   "source": [
    "Note that the outputs of `pvtol_noisy` are not the full set of states, but rather the states we can measure: $x$, $y$, and $\\theta$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce469b3-faa0-4bac-b9d4-02e4dae7a2da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function tlot the trajectory in xy coordinates\n",
    "def plot_results(t, x, u, fig=None):\n",
    "    # Set the size of the figure\n",
    "    if fig is None:\n",
    "        fig = plt.figure(figsize=(10, 6))\n",
    "\n",
    "    # Top plot: xy trajectory\n",
    "    plt.subplot(2, 1, 1)\n",
    "    plt.plot(x[0], x[1])\n",
    "    plt.xlabel('x [m]')\n",
    "    plt.ylabel('y [m]')\n",
    "    plt.axis('equal')\n",
    "\n",
    "    # Time traces of the state and input\n",
    "    plt.subplot(2, 4, 5)\n",
    "    plt.plot(t, x[1])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('y [m]')\n",
    "\n",
    "    plt.subplot(2, 4, 6)\n",
    "    plt.plot(t, x[2])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('theta [rad]')\n",
    "\n",
    "    plt.subplot(2, 4, 7)\n",
    "    plt.plot(t, u[0])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('$F_1$ [N]')\n",
    "\n",
    "    plt.subplot(2, 4, 8)\n",
    "    plt.plot(t, u[1])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('$F_2$ [N]')\n",
    "    plt.tight_layout()\n",
    "\n",
    "    return fig\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "081764e0",
   "metadata": {
    "id": "081764e0"
   },
   "source": [
    "## Estimator\n",
    "\n",
    "We start by designing an optimal estimator for the system.  We choose the noise intensities\n",
    "based on knowledge of the modeling errors, disturbances, and sensor characteristics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "778fb908",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disturbance and noise intensities\n",
    "Qv = np.diag([1e-2, 1e-2])\n",
    "Qw = np.array([[2e-4, 0, 1e-5], [0, 2e-4, 1e-5], [1e-5, 1e-5, 1e-4]])\n",
    "Qwinv = np.linalg.inv(Qw)\n",
    "\n",
    "# Initial state covariance\n",
    "P0 = np.eye(pvtol_noisy.nstates)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1Q55PHN1omJs",
   "metadata": {
    "id": "1Q55PHN1omJs"
   },
   "source": [
    "We will use a linear quadratic estimator (Kalman filter) to design an optimal estimator for the system. Recall that the `ct.lqe` function takes in a linear system as input, so we first linear our `pvtol_noisy` system around its equilibrium point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "WADb1-VcuR5t",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the equilibrium point corresponding to the origin\n",
    "xe, ue = ct.find_eqpt(\n",
    "    sys = pvtol_noisy,\n",
    "    x0 = np.zeros(pvtol_noisy.nstates),\n",
    "    u0 = np.zeros(pvtol_noisy.ninputs),\n",
    "    y0 = [0, 0, 0],\n",
    "    iu=range(2, pvtol_noisy.ninputs),\n",
    "    iy=[0, 1]\n",
    ")\n",
    "print(f\"{xe=}\")\n",
    "print(f\"{ue=}\")\n",
    "\n",
    "# Linearize system for Kalman filter\n",
    "pvtol_noisy_lin = pvtol_noisy.linearize(xe, ue)\n",
    "\n",
    "# Extract the linearization for use in LQR design\n",
    "A, B, C = pvtol_noisy_lin.A, pvtol_noisy_lin.B, pvtol_noisy_lin.C"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6E9s147Cpppr",
   "metadata": {
    "id": "6E9s147Cpppr"
   },
   "source": [
    "We want to define an estimator that takes in the measured states $x$, $y$, and $\\theta$, as well as applied inputs $F_1$ and $F_2$. As the estimator doesn't have any measurement of the noise/disturbances applied to the system, we will design our controller with only these inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "nvZHm0Ooqkj_",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use ct.lqe to create an L matrix, using only measured inputs F1 and F2\n",
    "L, Pf, _ = ct.lqe(A, B[:,:2], C, Qv, Qw)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "KXVetnCUrHvs",
   "metadata": {
    "id": "KXVetnCUrHvs"
   },
   "source": [
    "We now create our estimator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "M77vo5PgrIEv",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create standard (optimal) estimator update function\n",
    "def estimator_update(t, xhat, u, params):\n",
    "\n",
    "    # Extract the inputs to the estimator\n",
    "    y = u[0:3]                  # just grab the first three outputs\n",
    "    u_cmd = u[3:5]              # get the inputs that were applied as well\n",
    "\n",
    "    # Update the state estimate using PVTOL (non-noisy) dynamics\n",
    "    return _pvtol_update(t, xhat, u_cmd, params) - L @ (C @ xhat - y)\n",
    "\n",
    "# Create estimator\n",
    "estimator = ct.nlsys(\n",
    "    estimator_update, None,\n",
    "    name = 'Estimator',\n",
    "    states=pvtol_noisy.nstates,\n",
    "    inputs= pvtol_noisy.output_labels  \\\n",
    "        + pvtol_noisy.input_labels[0:2],\n",
    "    outputs=[f'xh{i}' for i in range(pvtol_noisy.nstates)],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1JOPx1TXrnr-",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(estimator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d8463d",
   "metadata": {
    "id": "46d8463d"
   },
   "source": [
    "## Gain scheduled controller\n",
    "\n",
    "We next design our (gain scheduled) controller for the system. Here, as in the case of the estimator, we will create the controller using the nominal PVTOL system, so that the applied inputs to the system are only $F_1$ and $F_2$. If we were to make a controller using the noisy PVTOL system, then the inputs applied via control action would include noise and disturbances, which is incorrect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e5fbef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the weights for the LQR problem\n",
    "Qx = np.diag([100, 10, (180/np.pi) / 5, 0, 0, 0])\n",
    "# Qx = np.diag([10, 100, (180/np.pi) / 5, 0, 0, 0])    # Try this out to see what changes\n",
    "Qu = np.diag([10, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5cc3cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct the array of gains and the gain scheduled controller\n",
    "import itertools\n",
    "import math\n",
    "\n",
    "# Set up points around which to linearize (control-0.9.3: must be 2D or greater)\n",
    "angles = np.linspace(-math.pi/3, math.pi/3, 10)\n",
    "speeds = np.linspace(-10, 10, 3)\n",
    "points = list(itertools.product(angles, speeds))\n",
    "\n",
    "# Compute the gains at each design point of angles and speeds\n",
    "gains = []\n",
    "\n",
    "# Iterate through points\n",
    "for point in points:\n",
    "\n",
    "    # Compute the state that we want to linearize about\n",
    "    xgs = xe.copy()\n",
    "    xgs[2], xgs[4] = point[0], point[1]\n",
    "\n",
    "    # Linearize the system and compute the LQR gains\n",
    "    linsys = pvtol_noisy.linearize(xgs, ue)\n",
    "    A = linsys.A\n",
    "    B = linsys.B[:,:2]\n",
    "    K, X, E = ct.lqr(A, B, Qx, Qu)\n",
    "    gains.append(K)\n",
    "\n",
    "# Construct the controller\n",
    "gs_ctrl, gs_clsys = ct.create_statefbk_iosystem(\n",
    "    sys = pvtol_nominal,\n",
    "    gain = (gains, points),\n",
    "    gainsched_indices=['xh2', 'xh4'],\n",
    "    estimator=estimator\n",
    ")\n",
    "\n",
    "print(gs_ctrl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd28a73",
   "metadata": {
    "id": "ecd28a73"
   },
   "source": [
    "## Trajectory generation\n",
    "\n",
    "Finally, we need to design the trajectory that we want to follow.  We consider a situation with state constraints that represent the specific experimental conditions for this system (at Caltech):\n",
    "* `ceiling`: The system has limited vertical travel, so we constrain the vertical position to lie between $-0.5$ and $2$ meters.\n",
    "* `nicolas`: When testing, we placed a person in between the initial and final position, and we need to avoid hitting him as we move from start to finish.\n",
    "\n",
    "The code below defines the initial conditions, final conditions, and constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eb12bfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the initial and final conditions\n",
    "x_delta = np.array([10, 0, 0, 0, 0, 0])\n",
    "x0, u0 = ct.find_eqpt(\n",
    "    sys = pvtol_nominal,\n",
    "    x0 = np.zeros(6),\n",
    "    u0 = np.zeros(2),\n",
    "    y0 = np.zeros(3),\n",
    "    iy=[0, 1]\n",
    ")\n",
    "xf, uf = ct.find_eqpt(\n",
    "    sys = pvtol_nominal,\n",
    "    x0 = x0 + x_delta,\n",
    "    u0 = u0,\n",
    "    y0 = (x0 + x_delta)[:3],\n",
    "    iy=[0, 1]\n",
    ")\n",
    "\n",
    "# Define the time horizon for the manuever\n",
    "Tf = 5\n",
    "timepts = np.linspace(0, Tf, 100, endpoint=False)\n",
    "\n",
    "# Create a constraint corresponding to the obstacle\n",
    "ceiling = (NonlinearConstraint, lambda x, u: x[1], [-0.5], [2])\n",
    "nicolas = (NonlinearConstraint,\n",
    "           lambda x, u: x[1] - (0.5 * sin(pi * x[0] / 10) - 0.1), [0], [1])\n",
    "\n",
    "# # Reset the nonlinear constraint to give some extra room\n",
    "# nicolas = (NonlinearConstraint,\n",
    "#            lambda x, u: x[1] - (0.8 * sin(pi * x[0] / 10) - 0.1), [0], [1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "610aa247",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Re-define the time horizon for the manuever\n",
    "Tf = 5\n",
    "timepts = np.linspace(0, Tf, 20, endpoint=False)\n",
    "\n",
    "# We provide a tent shape as an intial guess\n",
    "xm = (x0 + xf) / 2 + np.array([0, 0.5, 0, 0, 0, 0])\n",
    "tm = int(len(timepts)/2)\n",
    "# straight line from start to midpoint to end with nominal input\n",
    "tent = (\n",
    "    np.hstack([\n",
    "      np.array([x0 + (xm - x0) * t/(Tf/2) for t in timepts[0:tm]]).transpose(),\n",
    "      np.array([xm + (xf - xm) * t/(Tf/2) for t in timepts[0:tm]]).transpose()\n",
    "    ]),\n",
    "    u0\n",
    ")\n",
    "\n",
    "# terminal constraint\n",
    "term_constraints = opt.state_range_constraint(pvtol_nominal, xf, xf)\n",
    "\n",
    "# trajectory cost\n",
    "traj_cost = opt.quadratic_cost(pvtol_nominal, None, Qu, x0=xf, u0=uf)\n",
    "\n",
    "# find optimal trajectory\n",
    "start_time = time.process_time()\n",
    "traj = opt.solve_ocp(\n",
    "    sys = pvtol_nominal,\n",
    "    timepts = timepts,\n",
    "    initial_guess=tent,\n",
    "    X0=x0,\n",
    "    cost = traj_cost,\n",
    "    trajectory_constraints=[ceiling, nicolas],\n",
    "    terminal_constraints=term_constraints,\n",
    ")\n",
    "print(\"* Total time = %5g seconds\\n\" % (time.process_time() - start_time))\n",
    "\n",
    "# Create the desired trajectory\n",
    "xd, ud = traj.states, traj.inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e59ddc29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extend the trajectory to hold the final position for Tf seconds\n",
    "holdpts = np.arange(Tf, Tf + Tf, timepts[1]-timepts[0])\n",
    "xd = np.hstack([xd, np.outer(xf, np.ones_like(holdpts))])\n",
    "ud = np.hstack([ud, np.outer(uf, np.ones_like(holdpts))])\n",
    "timepts = np.hstack([timepts, holdpts])\n",
    "\n",
    "# Plot the desired trajectory\n",
    "plot_results(timepts, xd, ud)\n",
    "plt.suptitle('Desired Trajectory')\n",
    "\n",
    "# Add the constraints to the plot\n",
    "plt.subplot(2, 1, 1)\n",
    "\n",
    "plt.plot([0, 10], [2, 2], 'r--')\n",
    "plt.text(5, 1.8, 'Ceiling', ha='center')\n",
    "\n",
    "x_nic = np.linspace(0, 10, 50)\n",
    "y_nic = 0.5 * np.sin(pi * x_nic / 10) - 0.1\n",
    "plt.plot(x_nic, y_nic, 'r--')\n",
    "plt.text(5, 0, 'Nicolas Petit', ha='center')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affe55fa",
   "metadata": {
    "id": "affe55fa"
   },
   "source": [
    "## Final Control System Implementation\n",
    "\n",
    "We now put together the final control system and simulate it.  If you have named your inputs and outputs to each of the subsystems properly, the code below should connect everything up correctly.  If you get errors about inputs or outputs that are not connected to anything, check the names of your inputs and outputs in the various\n",
    "systems above and make sure everything lines up as it should."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50dff557",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the interconnected system\n",
    "clsys = ct.interconnect(\n",
    "    [pvtol_noisy, gs_ctrl, estimator],\n",
    "    inputs=gs_clsys.input_labels[:8] + pvtol_noisy.input_labels[2:],\n",
    "    outputs=pvtol_noisy.output_labels + pvtol_noisy.input_labels[:2]\n",
    ")\n",
    "print(clsys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f24e6f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate disturbance and noise vectors\n",
    "V = ct.white_noise(timepts, Qv)\n",
    "W = ct.white_noise(timepts, Qw)\n",
    "for i in range(V.shape[0]):\n",
    "    plt.subplot(2, 3, i+1)\n",
    "    plt.plot(timepts, V[i])\n",
    "    plt.ylabel(f'V[{i}]')\n",
    "\n",
    "for i in range(W.shape[0]):\n",
    "    plt.subplot(2, 3, i+4)\n",
    "    plt.plot(timepts, W[i])\n",
    "    plt.ylabel(f'W[{i}]')\n",
    "    plt.xlabel('Time $t$ [s]')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63091cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate the open loop system and plot the results (+ state trajectory)\n",
    "resp = ct.input_output_response(\n",
    "    sys = clsys,\n",
    "    T = timepts,\n",
    "    U = [xd, ud, V, W],\n",
    "    X0 = np.zeros(12))\n",
    "\n",
    "plot_results(resp.time, resp.outputs[0:3], resp.outputs[3:5])\n",
    "\n",
    "# Add the constraints to the plot\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot([0, 10], [1, 1], 'r--')\n",
    "x_nic = np.linspace(0, 10, 50)\n",
    "y_nic = 0.5 * np.sin(pi * x_nic / 10) - 0.1\n",
    "plt.plot(x_nic, y_nic, 'r--')\n",
    "plt.text(5, 0, 'Nicolas Petit', ha='center')\n",
    "plt.suptitle(\"Measured Trajectory\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89221230",
   "metadata": {
    "id": "89221230"
   },
   "source": [
    "We see that with the addition of disturbances and noise, we sometimes violate the constraint 'nicolas' (if your plot doesn't show an intersection with the bottom dashed curve, try regenerating the noise and running the simulation again).  This can be fixed by establishing a more conservative constraint (see commented out constraint in code block above)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f2e9776-0ba9-4295-9473-a17cb4854836",
   "metadata": {
    "id": "3f2e9776-0ba9-4295-9473-a17cb4854836"
   },
   "source": [
    "## Small signal analysis\n",
    "\n",
    "We next look at the properties of the system using the small signal (linearized) dynamics.  This analysis is useful to check the robustness and performance of the controller around trajectories and equilibrium points.\n",
    "\n",
    "We will carry out the analysis around the initial condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "JgZyPyMkcoOl",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Small signal analysis\n",
    "X0 = np.hstack([x0, x0])                # system state, estim state\n",
    "U0 = np.hstack([x0, u0, np.zeros(5)])   # xd, ud, dist, noise\n",
    "G = clsys.linearize(X0, U0)\n",
    "print(clsys)\n",
    "\n",
    "# Get input/output dictionaries: inp['sig'] = index for 'sig'\n",
    "inp = clsys.input_index\n",
    "out = clsys.output_index\n",
    "\n",
    "fig, axs = plt.subplots(2, 3, figsize=[9, 6])\n",
    "omega = np.logspace(-2, 2)\n",
    "\n",
    "# Complementary sensitivity\n",
    "G_x_xd = ct.tf(G[out['x'], inp['xd[0]']])\n",
    "G_y_yd = ct.tf(G[out['y'], inp['xd[1]']])\n",
    "ct.bode_plot(\n",
    "    [G_x_xd, G_y_yd], omega,\n",
    "    plot_phase=False, ax=np.array([[axs[0, 0]]]))\n",
    "axs[0, 0].legend(['F T_x', 'F T_y'])\n",
    "axs[0, 0].loglog([omega[0], omega[-1]], [1, 1], 'k', linewidth=0.5)\n",
    "axs[0, 0].set_title(\"From xd, yd\", fontsize=9)\n",
    "axs[0, 0].set_ylabel(\"To x, y\")\n",
    "axs[0, 0].set_xlabel(\"\")\n",
    "\n",
    "# Load (or input) sensitivity\n",
    "G_x_dx = ct.tf(G[out['x'], inp['Dx']])\n",
    "G_y_dy = ct.tf(G[out['y'], inp['Dy']])\n",
    "ct.bode_plot(\n",
    "    [G_x_dx, G_y_dy], omega,\n",
    "    plot_phase=False, ax=np.array([[axs[0, 1]]]))\n",
    "axs[0, 1].legend(['PS_x', 'PS_y'])\n",
    "axs[0, 1].loglog([omega[0], omega[-1]], [1, 1], 'k', linewidth=0.5)\n",
    "axs[0, 1].set_title(\"From Dx, Dy\", fontsize=9)\n",
    "axs[0, 1].set_xlabel(\"\")\n",
    "axs[0, 1].set_ylabel(\"\")\n",
    "\n",
    "# Sensitivity\n",
    "G_x_Nx = ct.tf(G[out['x'], inp['Nx']])\n",
    "G_y_Ny = ct.tf(G[out['y'], inp['Ny']])\n",
    "ct.bode_plot(\n",
    "    [G_x_Nx, G_y_Ny], omega,\n",
    "    plot_phase=False, ax=np.array([[axs[0, 2]]]))\n",
    "axs[0, 2].legend(['S_x', 'S_y'])\n",
    "axs[0, 2].set_title(\"From Nx, Ny\", fontsize=9)\n",
    "axs[0, 2].loglog([omega[0], omega[-1]], [1, 1], 'k', linewidth=0.5)\n",
    "axs[0, 2].set_xlabel(\"\")\n",
    "axs[0, 2].set_ylabel(\"\")\n",
    "\n",
    "# Noise (or output) sensitivity\n",
    "G_F1_xd = ct.tf(G[out['F1'], inp['xd[0]']])\n",
    "G_F2_yd = ct.tf(G[out['F2'], inp['xd[1]']])\n",
    "ct.bode_plot(\n",
    "    [G_F1_xd, G_F2_yd], omega,\n",
    "    plot_phase=False, ax=np.array([[axs[1, 0]]]))\n",
    "axs[1, 0].legend(['FCS_x', 'FCS_y'])\n",
    "axs[1, 0].loglog([omega[0], omega[-1]], [1, 1], 'k', linewidth=0.5)\n",
    "axs[1, 0].set_ylabel(\"To F1, F2\")\n",
    "\n",
    "G_F1_dx = ct.tf(G[out['F1'], inp['Dx']])\n",
    "G_F2_dy = ct.tf(G[out['F2'], inp['Dy']])\n",
    "ct.bode_plot(\n",
    "    [G_F1_dx, G_F2_dy], omega,\n",
    "    plot_phase=False, ax=np.array([[axs[1, 1]]]))\n",
    "axs[1, 1].legend(['~T_x', '~T_y'])\n",
    "axs[1, 1].loglog([omega[0], omega[-1]], [1, 1], 'k', linewidth=0.5)\n",
    "axs[1, 1].set_ylabel(\"\")\n",
    "\n",
    "# Sensitivity\n",
    "G_F1_Nx = ct.tf(G[out['F1'], inp['Nx']])\n",
    "G_F1_Ny = ct.tf(G[out['F1'], inp['Ny']])\n",
    "ct.bode_plot(\n",
    "    [G_F1_Nx, G_F1_Ny], omega,\n",
    "    plot_phase=False, ax=np.array([[axs[1, 2]]]))\n",
    "axs[1, 2].legend(['C S_x', 'C S_y'])\n",
    "axs[1, 2].loglog([omega[0], omega[-1]], [1, 1], 'k', linewidth=0.5)\n",
    "axs[1, 2].set_ylabel(\"\")\n",
    "\n",
    "plt.suptitle(\"Gang of Six for PVTOL\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "xfi1mXJTe3Gm",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve for the loop transfer function horizontal direction\n",
    "# S = 1 / (1 + L) => S + SL = 1 => L = (1 - S)/S\n",
    "Lx = (1 - G_x_Nx) / G_x_Nx; Lx.name = 'Lx'\n",
    "Ly = (1 - G_y_Ny) / G_y_Ny; Ly.name = 'Ly'\n",
    "\n",
    "# Create Nyquist plot\n",
    "ct.nyquist_plot([Lx, Ly], max_curve_magnitude=5, max_curve_offset=0.2);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "L7L6UZTn_Qtn",
   "metadata": {
    "id": "L7L6UZTn_Qtn"
   },
   "source": [
    "### Gain Margins of $L_x$, $L_y$\n",
    "\n",
    "We can zoom in on the plot to see the gain, phase, and stability margins:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3FX7YXrR2cuQ",
   "metadata": {},
   "outputs": [],
   "source": [
    "cplt = ct.nyquist_plot([Lx, Ly])\n",
    "lower_upper_bound = 1.1\n",
    "cplt.axes[0, 0].set_xlim([-lower_upper_bound, lower_upper_bound])\n",
    "cplt.axes[0, 0].set_ylim([-lower_upper_bound, lower_upper_bound])\n",
    "cplt.axes[0, 0].set_aspect('equal')\n",
    "\n",
    "# Gain margin for Lx\n",
    "neg1overgm_x = -0.67  # vary this manually to find intersection with curve\n",
    "color = cplt.lines[0, 0][0].get_color()\n",
    "plt.plot(neg1overgm_x, 0, color=color, marker='o', fillstyle='none')\n",
    "gm_x = -1/neg1overgm_x\n",
    "\n",
    "# Gain margin for Ly\n",
    "neg1overgm_y = -0.32  # vary this manually to find intersection with curve\n",
    "color = cplt.lines[1, 0][0].get_color()\n",
    "plt.plot(neg1overgm_y, 0, color=color, marker='o', fillstyle='none')\n",
    "gm_y = -1/neg1overgm_y\n",
    "\n",
    "print('Margins obtained visually:')\n",
    "print('Gain margin of Lx: '+str(gm_x))\n",
    "print('Gain margin of Ly: '+str(gm_y))\n",
    "print('\\n')\n",
    "\n",
    "# get gain margin computationally\n",
    "gm_xc, pm_xc, wpc_xc, wgc_xc = ct.margin(Lx)\n",
    "gm_yc, pm_yc, wpc_yc, wgc_yc = ct.margin(Ly)\n",
    "\n",
    "print('Margins obtained computationally:')\n",
    "print('Gain margin of Lx: '+str(gm_xc))\n",
    "print('Gain margin of Ly: '+str(gm_yc))\n",
    "\n",
    "print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "VnrVNvhz_Zi2",
   "metadata": {
    "id": "VnrVNvhz_Zi2"
   },
   "source": [
    "### Phase Margins of $L_x$, $L_y$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "zKb_o9ZN_ffF",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add customizations to Nyquist plot\n",
    "cplt = ct.nyquist_plot(\n",
    "    [Lx, Ly], max_curve_magnitude=5, max_curve_offset=0.2,\n",
    "    unit_circle=True)\n",
    "lower_upper_bound = 2\n",
    "cplt.axes[0, 0].set_xlim([-lower_upper_bound, lower_upper_bound])\n",
    "cplt.axes[0, 0].set_ylim([-lower_upper_bound, lower_upper_bound])\n",
    "cplt.axes[0, 0].set_aspect('equal')\n",
    "\n",
    "# Phase margin of Lx:\n",
    "th_pm_x = 0.14*np.pi\n",
    "th_plt_x = np.pi + th_pm_x\n",
    "color = cplt.lines[0, 0][0].get_color()\n",
    "plt.plot(np.cos(th_plt_x), np.sin(th_plt_x), color=color, marker='o')\n",
    "\n",
    "# Phase margin of Ly\n",
    "th_pm_y = 0.19*np.pi\n",
    "th_plt_y = np.pi + th_pm_y\n",
    "color = cplt.lines[1, 0][0].get_color()\n",
    "plt.plot(np.cos(th_plt_y), np.sin(th_plt_y), color=color, marker='o')\n",
    "\n",
    "print('Margins obtained visually:')\n",
    "print('Phase margin: '+str(float(th_pm_x)))\n",
    "print('Phase margin: '+str(float(th_pm_y)))\n",
    "print('\\n')\n",
    "\n",
    "# get margin computationally\n",
    "gm_xc, pm_xc, wpc_xc, wgc_xc = ct.margin(Lx)\n",
    "gm_yc, pm_yc, wpc_yc, wgc_yc = ct.margin(Ly)\n",
    "\n",
    "print('Margins obtained computationally:')\n",
    "print('Phase margin of Lx: '+str(np.deg2rad(pm_xc)))\n",
    "print('Phase margin of Ly: '+str(np.deg2rad(pm_yc)))\n",
    "\n",
    "print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dF0BIq5BDXII",
   "metadata": {
    "id": "dF0BIq5BDXII"
   },
   "source": [
    "### Stability Margins of $L_x$, $L_y$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "XQPB_h6Y1cAW",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add customizations to Nyquist plot\n",
    "cplt = ct.nyquist_plot([Lx, Ly], max_curve_magnitude=5, max_curve_offset=0.2)\n",
    "lower_upper_bound = 2\n",
    "cplt.axes[0, 0].set_xlim([-lower_upper_bound, lower_upper_bound])\n",
    "cplt.axes[0, 0].set_ylim([-lower_upper_bound, lower_upper_bound])\n",
    "cplt.axes[0, 0].set_aspect('equal')\n",
    "\n",
    "# Stability margin:\n",
    "sm_x = 0.3  # vary this manually to find min which intersects\n",
    "color = cplt.lines[0, 0][0].get_color()\n",
    "sm_circle = plt.Circle((-1, 0), sm_x, color=color, fill=False, ls=':')\n",
    "cplt.axes[0, 0].add_patch(sm_circle)\n",
    "\n",
    "sm_y = 0.5  # vary this manually to find min which intersects\n",
    "color = cplt.lines[1, 0][0].get_color()\n",
    "sm_circle = plt.Circle((-1, 0), sm_y, color=color, fill=False, ls=':')\n",
    "cplt.axes[0, 0].add_patch(sm_circle)\n",
    "\n",
    "print('Margins obtained visually:')\n",
    "print('* Stability margin of Lx: '+str(sm_x))\n",
    "print('* Stability margin of Ly: '+str(sm_y))\n",
    "\n",
    "# Compute the stability margin computationally\n",
    "print('')  # blank line\n",
    "print('Margins obtained computationally:')\n",
    "resp = ct.frequency_response(1 + Lx)\n",
    "sm = np.min(resp.magnitude)\n",
    "wsm = resp.omega[np.argmin(resp.magnitude)]\n",
    "\n",
    "print(f\"* Stability margin of Lx = {sm:2.2g} (at {wsm:.2g} rad/s)\")\n",
    "resp = ct.frequency_response(1 + Ly)\n",
    "sm = np.min(resp.magnitude)\n",
    "wsm = resp.omega[np.argmin(resp.magnitude)]\n",
    "print(f\"* Stability margin of Ly = {sm:2.2g} (at {wsm:.2g} rad/s)\")\n",
    "print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "boAjWk56GXYZ",
   "metadata": {
    "id": "boAjWk56GXYZ"
   },
   "source": [
    "We see that the frequencies at which the stability margins are found corresponds to the peak of the magnitude of the sensitivity functions for $L_x$ and $L_y$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "JkbMn8pif7Ub",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Confirm stability using Nyquist criterion\n",
    "nyqresp_x = ct.nyquist_response(Lx)\n",
    "nyqresp_y = ct.nyquist_response(Ly)\n",
    "\n",
    "print(\"Nx =\", nyqresp_x.count, \"; Px =\", np.sum(np.real(Lx.poles()) > 0))\n",
    "print(\"Ny =\", nyqresp_y.count, \"; Py =\", np.sum(np.real(Ly.poles()) > 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d038db9-f671-4f0f-82db-51096e8272b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Take a look at the locations of the poles\n",
    "np.real(Ly.poles())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dd57510-4b03-4c0a-90ae-35011f90c41b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See what happened in the contour\n",
    "plt.plot(np.real(nyqresp_y.contour), np.imag(nyqresp_y.contour))\n",
    "plt.axis([-1e-4, 4e-4, 0, 4e-4])\n",
    "plt.title(\"Zoom on D-contour\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7b9a2f9-f40f-4090-ae69-6bf53fea54a9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
